import pyqtgraph as pg

def plot_graph_from_data(data: dict):
    """
    生成GraphItem。
    :param data: 包含节点坐标、边的连接关系和风格的数据。
    """
    def hex_to_rgba(hex_str, default_alpha=255):
        """
        将十六进制颜色字符串（如 '#FF0000' 或 '#FF000080'）转换为 [r, g, b, a] 整数列表。
        支持格式：
        - '#RRGGBB' -> [255, 0, 0, 255]
        - '#RRGGBBAA' -> [255, 0, 0, 128]
        :param hex_str: 16进制字符串。
        :param default_alpha: 默认透明度, defaults to 255
        :raises ValueError: 无效格式，必须 # 开头的6或者8位16进制数。
        :return: 返回RGBA格式列表。
        """
        hex_str = hex_str.lstrip('#')
        r = int(hex_str[0: 2], 16)
        g = int(hex_str[2: 4], 16)
        b = int(hex_str[4: 6], 16)
        if len(hex_str) == 6:
            return [r, g, b, default_alpha]
        elif len(hex_str) == 8:
            a = int(hex_str[6: 8], 16)
            return [r, g, b, a]
        else:
            raise ValueError(f"Invalid hex color format: {hex_str}")
    
    import numpy as np
    pos = np.array(data['pos'])
    adj = np.array(data['adj'])

    default_node_style = data.get('default_node_style', {
        'color': '#FF0000',
        'border_color': '#AA0000',
        'size': 10,
        'shape': 'o'
    })

    default_edge_style = data.get('default_edge_style', {
        'color': '#0000FF',
        'width': 2
    })

    node_styles = data.get('node_styles', {})  # key: "0", "1", ...
    edge_styles = data.get('edge_styles', {})  # key: "0", "4", ...

    num_nodes = len(pos)
    num_edges = len(adj)

    # ===== 1. 准备节点数据 =====
    sizes = []
    shapes = []
    node_colors = []
    border_colors = []
    symbol_pens = []
    symbol_brushes = []

    for i in range(num_nodes):
        node_idx_str = str(i)
        style = node_styles.get(node_idx_str, default_node_style)

        color = style.get('color', default_node_style['color'])
        border_color = style.get('border_color', color)
        size = style.get('size', default_node_style['size'])
        shape = style.get('shape', default_node_style['shape'])

        sizes.append(size)
        shapes.append(shape)
        node_colors.append(color)
        border_colors.append(border_color)

        pen = pg.mkPen(color=border_color, width=2)
        brush = pg.mkBrush(color=color)
        symbol_pens.append(pen)
        symbol_brushes.append(brush)

    # ===== 2. 准备边数据：构建 pen_ndarray = [ (r, g, b, a, width), ... ] =====
    pen_data = []
    for edge_i in range(num_edges):  # 遍历每条边，edge_i 是 0, 1, 2...
        edge_idx_str = str(edge_i)  # 边的索引是字符串，如 "0", "1", ...
        style = edge_styles.get(edge_idx_str, default_edge_style)

        color = style.get('color', default_edge_style['color'])
        width = style.get('width', default_edge_style['width'])
        # 将十六进制颜色如 '#FF0000' 转为 [255, 0, 0, 255]
        rgba = hex_to_rgba(color)
        pen_data.append(tuple(rgba + [width]))  # [r, g, b, a, width]

    pen_ndarray = np.array(pen_data, 
                           dtype=[('red', np.ubyte), ('green', np.ubyte), ('blue', np.ubyte), ('alpha', np.ubyte), 
                            ('width', float)])  # shape: (num_edges, 5)
    # ===== 3. 构建 GraphItem 并设置数据 =====
    graph_item = pg.GraphItem()
    graph_item.setData(
        pos=pos,
        adj=adj,
        size=sizes,
        symbol=shapes,
        pen=pen_ndarray, 
        symbolPen=symbol_pens,
        symbolBrush=symbol_brushes,
        pxMode=False
    )
    return graph_item

if __name__ == "__main__":
    data = {
        "pos": [
            [0, 0],
            [100, 0],
            [0, 100],
            [100, 100],
            [50, 50],
            [150, 50]
        ],
        "adj": [
            [0, 1],
            [1, 3],
            [3, 2],
            [2, 0],
            [1, 4],
            [3, 4]
        ],
        "node_styles": {
            "1": {"color": "#00FF00", "size": 20, "shape": "s"},
            "3": {"color": "#FFFF00", "size": 15, "shape": "t"}
        },
        "edge_styles": {
            "0": {"color": "#FF0000", "width": 3},
            "4": {"color": "#0000FF", "width": 4}
        },
        "default_node_style": {
            "color": "#FF0000",
            "size": 10,
            "shape": "o"
        },
        "default_edge_style": {
            "color": "#0000FF",
            "width": 2
        }
    }
    # 示例调用（禁用鼠标，自定义样式）。
    graph = plot_graph_from_data(data)
    # 创建绘图窗口。
    win = pg.GraphicsLayoutWidget()
    win.setWindowTitle("Styled Graph")
    win.resize(800, 600)
    # 添加绘图区域。
    view: pg.ViewBox = win.addViewBox()
    view.setAspectLocked(True)  # 保持纵横比。
    view.addItem(graph)
    win.show()
    pg.exec()